import collections
import os
import traceback

import numpy as np
import pandas as pd
from loguru import logger

from models.global_config import glm
from mylib import gen_excel
from models.stock_model import StockNumber, DayInfo


def get_all_stock(txt_name, trade_date, N, stocks, log_dir):
    df = pd.read_csv('all.csv')
    all_cnt = 0
    all_down_cnt = 0
    all_down_stock_dict = collections.defaultdict(list)

    indus_dict = collections.defaultdict(int)
    indus_all_dict = collections.defaultdict(int)
    indus_stock_dict = collections.defaultdict(list)
    indus_all_stock_dict = collections.defaultdict(list)
    for row in df.index:
        if stocks is not None and df.loc[row]['ts_code'] not in stocks:
            continue
        sn = StockNumber(df.loc[row])
        if '退' in sn.name:
            continue
        if 'ST' in sn.name:
            continue
        if not str(sn.ts_code).startswith('60') \
                and not str(sn.ts_code).startswith('30') \
                and not str(sn.ts_code).startswith('00'):
            continue
        all_cnt += 1
        indus_all_dict[txt_name] += 1
        indus_all_stock_dict[txt_name].append(sn.name)
        if analysis_stock(trade_date, N, sn, log_dir):
            all_down_stock_dict[txt_name].append(sn.name)
            all_down_cnt += 1
            indus_stock_dict[txt_name].append(sn.name)
            indus_dict[txt_name] += 1
            logger.info(f'{trade_date} cnt = {indus_dict[txt_name]}')

    with open(f'result/{all_cnt}_{txt_name}.csv', 'a+') as fw:
        msg = f'{trade_date},{all_down_cnt},{all_cnt},{list(all_down_stock_dict.items())}'
        logger.success(msg)
        fw.write(msg)
        fw.write('\n')

    # for k, v in indus_dict.items():
    #     with open(f'result/trade_date_cnt_{txt_name}.csv', 'a+') as fw:
    #         msg = f'{trade_date},{k}, {v}, {indus_all_dict.get(k)},' \
    #               f' {indus_stock_dict.get(k)}, {indus_all_stock_dict.get(k)}'
    #         logger.success(msg)
    #         fw.write(msg)
    #         fw.write('\n')


def analysis_stock(trade_date, N, sn, log_dir):
    csv_path = f'stocks/{sn.ts_code}.csv'
    df2 = pd.read_csv(csv_path)
    q10 = []
    q11 = []
    all_di = []
    parr = 0
    today_di = None
    cnt = 0
    pct_arr = []
    arr_high = []
    arr_low = []
    date_arr = []
    price_arr = []
    delta_days = []
    delta_d = 0
    price_arr_high = []
    price_arr_low = []
    for row2 in df2.index:
        delta_d += 1
        di = DayInfo(sn, df2.loc[row2])
        di.name = sn.name
        di.ts_code = sn.ts_code
        di.industry = sn.industry
        # di.trade_date = df2.loc[row2]['trade_date']
        if int(di.trade_date) > int(trade_date):
            continue
        # di.open = df2.loc[row2]['open']
        # di.high = df2.loc[row2]['high']
        price_arr_high.append(di.high)
        # di.low = df2.loc[row2]['low']
        price_arr_low.append(di.low)
        # di.close = float(df2.loc[row2]['close'])
        if cnt > 40:
            break
        # di.pre_close = df2.loc[row2]['pre_close']
        # di.change = df2.loc[row2]['change']
        # di.pct_chg = float(df2.loc[row2]['pct_chg'])
        # di.vol = df2.loc[row2]['vol']
        # di.amount = df2.loc[row2]['amount']
        if len(q10) == N * 2 + 1:
            ad = all_di[N]
            if max(q10) == q10[N]:
                delta_days.append(-delta_d) if delta_days else delta_days.append(-(delta_d - 6))
                delta_d = 0
                pct = round((parr - q10[N]) / q10[N] * 100, 2)
                pct_arr.append(pct)
                if pct > 0:
                    arr_high.append(pct)
                else:
                    arr_low.append(pct)
                price_arr.append(q10[N])
                date_arr.append(ad.trade_date)
                parr = q10[N]
                cnt += 1
            elif min(q11) == q11[N]:
                pct = round((parr - q11[N]) / q11[N] * 100, 2)
                pct_arr.append(pct)
                if pct > 0:
                    delta_days.append(delta_d) if delta_days else delta_days.append(delta_d - 6)
                    arr_high.append(pct)
                else:
                    delta_days.append(-delta_d) if delta_days else delta_days.append(-(delta_d - 6))
                    arr_low.append(pct)
                delta_d = 0
                price_arr.append(q11[N])
                date_arr.append(ad.trade_date)
                parr = q11[N]
                cnt += 1
            all_di.pop(0)
            q10.pop(0)
            q11.pop(0)
        all_di.append(di)
        q10.append(di.high)
        q11.append(di.low)
        if today_di is None:
            parr = di.low
            today_di = di
            price_arr.append(di.low)
            date_arr.append(di.trade_date)
    # cal
    median_high = round(np.median(arr_high), 2)
    median_low = round(np.median(arr_low), 2)
    link_code_arr = sn.ts_code.split('.')
    link_code = f'{link_code_arr[1]}{link_code_arr[0]}'
    hyperlink = f'"https://xueqiu.com/S/{link_code}"'
    hyperlink2 = f'= HYPERLINK("https://xueqiu.com/S/{link_code}"'
    date_arr_hyp = [f'= HYPERLINK({hyperlink},{d})' for d in date_arr[:-1]]
    data_dict = {
        'date': date_arr_hyp,
        'price': price_arr[:-1],
        'delta_days': delta_days,
        'all_pct': pct_arr,
        'median_high': [median_high] * len(pct_arr),
        'median_low': [median_low] * len(pct_arr),
    }
    file_path = f'{log_dir}/{sn.ts_code}_{sn.name}_{today_di.trade_date}.xlsx'
    if pct_arr[0] > 0:
        max_p = max(price_arr_high[0:abs(delta_days[0])])
        pre_price = round(max_p * (100 + median_low) / 100, 2)
        desc = f' {max_p} ↓ -{median_low}% → {pre_price}'
    else:
        min_p = min(price_arr_low[0:abs(delta_days[0])])
        pre_price = round(min_p * (100 + median_high) / 100, 2)
        desc = f' {min_p} ↑ {median_high}% → {pre_price}'
    title_arr_x = [
        f'arr_low({len(arr_low)}) = {np.sort(arr_low)}',
        f'arr_hight({len(arr_high)}) = {np.sort(arr_high)}',
    ]
    chart_x_title = '\n'.join(title_arr_x)
    title_arr_y = [
        f'median_hight = {median_high}%',
        f'median_low = {median_low}%'
    ]
    chart_y_title = '\n'.join(title_arr_y)
    # gen_excel.export_with_chart(p_data=data_dict,
    #                             p_title=f'{sn.ts_code}_{sn.name} {desc}',
    #                             p_x_title=chart_x_title,
    #                             p_y_title=chart_y_title,
    #                             p_file_path=file_path)
    if pct_arr[0] < median_low:
        # max_down_msg = f'{today_di.trade_date},{sn.industry},{hyperlink2},{sn.name},{sn.industry},{pct_arr[0]},小于平均跌幅,{median_low}'
        # logger.info(max_down_msg)
        # with open(f'result/cal_point2_{sn.industry}.csv', 'a+') as fw:
        #     fw.write(max_down_msg)
        #     fw.write('\n')
        return True
    return False


def run(txt_name, stocks=None, bkd=0, N=6):
    log_dir = 'main_name_for2'
    if not os.path.exists(log_dir):
        os.mkdir(log_dir)
    logger.info('main_name_for2 start')
    for idx, trade_date in enumerate(glm.all_trade_date_arr):
        if idx > bkd:
            break
        logger.info(f'cal {trade_date}')
        try:
            get_all_stock(txt_name, trade_date, N, stocks, log_dir)
        except Exception as e:
            logger.error(traceback.format_exc())
            logger.error(e)
    logger.info('生成excel运行完成')
